import argparse
import builtins
import json
import os
from wechat_ocr.ocr_manager import OcrManager, OCR_MAX_TASK_ID
current_path = os.getcwd()
configPath = os.path.join(current_path, 'config.json')
with open(configPath, 'r', encoding='utf-8') as f:
    config = json.load(f)

wechat_ocr_dir = config['wechat_ocr_dir']
wechat_dir = config['wechat_dir']

result = {}
# 1. 屏蔽所有 print
_builtin_print = builtins.print
builtins.print = lambda *_, **__: None   # 空操作
def ocr_result_callback(img_path: str, results: dict):
    global result
    result_file = img_path + ".json"
    print(f"识别成功，img_path: {img_path}, result_file: {result_file}")
    with open(result_file, 'w', encoding='utf-8') as f:
       f.write(json.dumps(results, ensure_ascii=False, indent=2))
    result[img_path] = results
def main():
    parser = argparse.ArgumentParser(description="微信OCR识别图片")
    parser.add_argument("image_paths", nargs="+", help="一个或多个图片路径")
    args = parser.parse_args()
    image_paths = args.image_paths
    ocr_manager = OcrManager(wechat_dir)
    # 设置WeChatOcr目录
    ocr_manager.SetExePath(wechat_ocr_dir)
    # 设置微信所在路径
    ocr_manager.SetUsrLibDir(wechat_dir)
    # 设置ocr识别结果的回调函数
    ocr_manager.SetOcrResultCallback(ocr_result_callback)
    # 启动ocr服务
    ocr_manager.StartWeChatOCR()
    for image_path in image_paths:
        if not os.path.isfile(image_path):
            continue
        ocr_manager.DoOCRTask(image_path)
    while ocr_manager.m_task_id.qsize() != OCR_MAX_TASK_ID:
        pass
    # 识别输出结果
    ocr_manager.KillWeChatOCR()
    # 2. 恢复 print，只输出最终结果
    builtins.print = _builtin_print
    print(json.dumps(result, ensure_ascii=False))
    return json.dumps(result, ensure_ascii=False)
main()